package com.castlzl.random.core.leaf.zookeeper;

import com.castlzl.random.core.leaf.LeafSegmentIdGenerator.Segment;
import com.castlzl.random.core.leaf.LeafSegmentService;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.recipes.atomic.AtomicValue;
import org.apache.curator.framework.recipes.atomic.DistributedAtomicLong;
import org.apache.curator.framework.recipes.atomic.PromotedToLock;
import org.apache.curator.retry.RetryNTimes;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;

import java.nio.charset.StandardCharsets;

/**
 * ZooKeeper实现的号段分配服务
 * 使用ZooKeeper的分布式原子计数器实现号段分配
 * @author leiziliang
 */
@Slf4j
public class ZookeeperSegmentServiceImpl implements LeafSegmentService {

    private final CuratorFramework client;
    private final long step;
    private final String rootPath;
    private static final String PREFIX = "/";

    public ZookeeperSegmentServiceImpl(CuratorFramework client, long step, String pathPrefix) {
        this.client = client;
        this.step = step;
        if (!pathPrefix.startsWith(PREFIX)) {
            this.rootPath = PREFIX + pathPrefix;
        } else {
            this.rootPath = pathPrefix;
        }
        initZookeeperPath();
    }
    
    private void initZookeeperPath() {
        try {
            // 确保根路径存在
            Stat stat = client.checkExists().forPath(rootPath);
            if (stat == null) {
                log.debug("Creating zookeeper root path: {}", rootPath);
                client.create()
                      .creatingParentsIfNeeded()
                      .withMode(CreateMode.PERSISTENT)
                      .forPath(rootPath);
                log.info("Successfully created zookeeper root path: {}", rootPath);
            } else {
                log.debug("Zookeeper root path already exists: {}", rootPath);
            }
            
            // 确保业务标签目录存在
            String bizTagsPath = rootPath + "/biztags";
            if (client.checkExists().forPath(bizTagsPath) == null) {
                client.create()
                      .creatingParentsIfNeeded()
                      .withMode(CreateMode.PERSISTENT)
                      .forPath(bizTagsPath);
                log.info("Successfully created biztags path: {}", bizTagsPath);
            }
        } catch (Exception e) {
            log.error("Failed to init zookeeper path: {}", rootPath, e);
            throw new RuntimeException("Failed to init zookeeper path: " + rootPath, e);
        }
    }

    @Override
    public Segment getNextSegment(String bizTag) throws Exception {
        String bizPath = rootPath + "/biztags/" + bizTag;
        
        // 确保业务节点存在
        try {
            Stat stat = client.checkExists().forPath(bizPath);
            if (stat == null) {
                try {
                    log.info("Creating new bizTag node: {} with initial value: 0", bizPath);
                    client.create()
                          .creatingParentsIfNeeded()
                          .withMode(CreateMode.PERSISTENT)
                          .forPath(bizPath, "0".getBytes(StandardCharsets.UTF_8));
                } catch (Exception e) {
                    // 忽略节点已存在的异常，可能是并发创建导致的
                    log.debug("Node already exists (concurrent creation): {}", bizPath);
                }
            }
        } catch (Exception e) {
            log.error("Error checking/creating bizTag node: {}", bizPath, e);
            throw new RuntimeException("Failed to ensure bizTag node exists: " + bizPath, e);
        }

        // 使用分布式原子计数器
        DistributedAtomicLong count = new DistributedAtomicLong(
            client,
            bizPath,
            new RetryNTimes(10, 30),
            PromotedToLock.builder()
                .lockPath(rootPath + "/locks/" + bizTag)
                .retryPolicy(new RetryNTimes(10, 30))
                .build()
        );

        // 增加并获取当前值
        AtomicValue<Long> value = count.add(step);
        if (!value.succeeded()) {
            log.error("Failed to increment counter for bizTag: {}", bizTag);
            throw new RuntimeException("Failed to get segment from ZooKeeper for bizTag: " + bizTag);
        }

        long current = value.postValue();
        long min = current - step + 1;
        long max = current;
        
        log.debug("Generated new segment for bizTag: {}, min: {}, max: {}", bizTag, min, max);
        return new Segment(min, max);
    }
}